import torch.nn as nn
from options import HiDDenConfiguration
from model.conv_bn_relu import ConvBNRelu

class Decoder(nn.Module):
    """
    Decoder module for extracting 1BPP watermark from an image.
    """

    def __init__(self, config: HiDDenConfiguration):
        super(Decoder, self).__init__()
        self.channels = config.decoder_channels

        layers = [ConvBNRelu(3, self.channels)]
        for _ in range(config.decoder_blocks - 1):
            layers.append(ConvBNRelu(self.channels, self.channels))

        # 输出 1 个通道（1BPP）
        layers.append(ConvBNRelu(self.channels, 1))  

        self.layers = nn.Sequential(*layers)

    def forward(self, image_with_wm):
        x = self.layers(image_with_wm)
        return x  # [batch_size, 1, H, W]
